from transformers import BertTokenizer

# 设置具体包含 config.json的目录
model_dir = r"D:\AI\model\bert-base-chinese"  #这是一个文本分类模型

# 加载模型和分词器
# model = BertForSequenceClassification.from_pretrained(model_dir)
tokenizer = BertTokenizer.from_pretrained(model_dir)

# print(tokenizer)
# exit()

sents = ["我今天很开心", "我昨天很生气"]

# 批量编码句子
out = tokenizer.batch_encode_plus(
    batch_text_or_text_pairs = [sents[0], sents[1]],
    add_special_tokens=True,
    #当句子长度大于max_length时，截断
    truncation=True,
    #填充，使其长度等于max_length
    max_length=30,
    #可取值tf,pt,np,默认为返回list
    return_tensors=None,
    #返回attention_mask
    return_attention_mask=True,
    #返回输入的token type id
    return_token_type_ids=True,
    #返回speical_tokens_mask, 特殊符号标识
    return_special_tokens_mask=True,
    #返回offset_mapping 表示每一个词的起止位置，这个参数值能BertTokenizerFast使用
    # return_offsets_mapping=True,
    #返回length 表示长度
    return_length=True,
    )
#input_ids 就是编码后的词
#token_type_ids 就是编码后的词的类别,第一个句子和特殊符号的位置是0，第二个句子的位置是1
#special_token_mask 就是特殊符号的标识,1表示特殊符号，0表示正常
#attention_mask pad的职位是0，其他位置是1
#length 返回句子的长度

print(out)
for k,v in out.items():
    print(k,v)

print(tokenizer.decode(out['input_ids'][0]), tokenizer.decode(out['input_ids'][1]))




# 获取字典
vocab = tokenizer.get_vocab()
# print(vocab)

# 判断字典中是否包含某个词
print("阳" in vocab)
print("阳光" in vocab)
print(len(vocab))
#添加新词
tokenizer.add_tokens(["阳光","大地"])
vocab = tokenizer.get_vocab()
print(len(vocab))
print("阳光" in vocab)

#添加新的特殊符号
tokenizer.add_special_tokens({"additional_special_tokens":["<sun>","<earth>"]})
tokenizer.add_special_tokens({"eos_token":"[EOS]"})


#编码新句子
out = tokenizer.encode(text="阳光照在大地上[EOS]", text_pair=None, truncation=True, padding="max_length", max_length=10, add_special_tokens=True, return_tensors=None)

print( out)

# 解码为原来字符串
print(tokenizer.decode(out))


